# Copyright (c) 2023 Presto Labs Pte. Ltd.
# Author: leon

import datetime
import logging
from typing import List
import sys

import tornado.httpserver
import tornado.websocket
import tornado.ioloop
import tornado.web
from absl import flags

from coin.flow.subscriber import from_queue
from coin.base.datetime_util import parse_end_time
from coin.exchange.base.ws import writer
from coin.exchange.base.ws.subscriber import RotatingSubscriber
from coin.exchange.uniswap_v3_arbitrum import config
from coin.exchange.uniswap_v3.rest.subscriber import UniswapV3Subscriber
from coin.exchange.uniswap_base.constants import PRODUCTION_RPC_URL

FLAGS = flags.FLAGS


class BaseSubscriberFactory(object):
    def __init__(self, group_id, worker_id, logger=None, group_symbol=None):
        self._group_id = group_id
        self._worker_id = worker_id
        self._logger = logger
        self._group_symbol = group_symbol or config.group_symbol

    def init_subscriber(self, queue, ioloop, close_callback):
        sub = UniswapV3Subscriber(queue,
                                  self._group_id,
                                  self._worker_id,
                                  self._group_symbol[self._group_id],
                                  provider=PRODUCTION_RPC_URL['arbitrum'],
                                  blockchain_name='arbitrum',
                                  close_callback=close_callback,
                                  logger=self._logger)
        sub.start(ioloop)
        return sub


class WsHandler(tornado.websocket.WebSocketHandler):
    clients = []

    def open(self):
        logging.info('New connection.')
        self.clients.append(self)

    def on_message(self, message):
        logging.info('Message received: %s' % message)

    def on_close(self):
        logging.info('Connection closed.')
        self.clients.remove(self)

    @classmethod
    def write_to_clients(cls, data):
        for client in cls.clients:
            client.write_message(data)


class RawFeedHandler(object):
    _clients: List[WsHandler]

    def __init__(self, clients: List[WsHandler], worker_id: int):
        self._clients = clients
        self._worker_id = worker_id

    def on_book(self, record, queue_data, topic_data):
        for client in self._clients:
            to_publish = {
                "timestamp": record.timestamp,
                "topic": "depth",
                "worker_id": self._worker_id,
                "data": record.data
            }
            try:
                client.write_message(to_publish)
            except tornado.websocket.WebSocketClosedError:
                self._clients.remove(client)

    def on_trade(self, record, queue_data, topic_data):
        for client in self._clients:
            client.write_message(record)


def main(argv):
    assert len(argv) >= 2, '<worker_id> is missing.'
    worker_id = argv[1].strip()
    if len(argv) == 3:
        groups = sorted(set([c for c in argv[2]]))
    else:
        groups = sorted(config.group_symbol.keys())
    logging.info('Subscribing groups: %s' % ', '.join(groups))

    unit_td = datetime.timedelta(hours=1)
    if int(worker_id) % 2 == 1:
        base_offset_td = datetime.timedelta(minutes=0)
        end_time = parse_end_time('23:59:00')
    else:
        base_offset_td = datetime.timedelta(minutes=30)
        end_time = parse_end_time('00:29:00')

    subscribers = []

    def init(queue, ioloop):
        for group_id in groups:
            group_idx = ord(group_id) - ord('a')
            offset_td = base_offset_td + datetime.timedelta(seconds=5 * group_idx)

            logger = logging.getLogger('uniswapv3-arbitrum_rest_subscriber_%s%s' % (group_id, worker_id))
            factory = BaseSubscriberFactory(group_id, worker_id, logger)
            sub = RotatingSubscriber(factory.init_subscriber, queue, unit_td, offset_td, logger=logger)
            sub.start(ioloop, datetime.timedelta(seconds=5 * group_idx))
            subscribers.append(sub)

            flow_sub = from_queue(queue)
            raw_feed_handler = RawFeedHandler(WsHandler.clients, int(worker_id))
            flow_sub.subscribe('uniswapv3-arbitrum_rest_depth_*%s' % worker_id, raw_feed_handler.on_book)
            flow_sub.subscribe('uniswapv3-arbitrum_rest_trade_*%s' % worker_id, raw_feed_handler.on_trade)

    def teardown():
        for s in subscribers:
            s.stop()

    application = tornado.web.Application([
        (r'/ws', WsHandler),
    ])
    http_server = tornado.httpserver.HTTPServer(application)
    http_server.listen(9444 + int(worker_id))

    writer.run2(init, teardown, end_time=end_time)


if __name__ == '__main__':
    logging.basicConfig(level='INFO', format='%(levelname)s %(asctime)s %(name)s] %(message)s')
    sys.exit(main(sys.argv))
